{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2025-03-25T02:16:42.731213Z",
     "start_time": "2025-03-25T02:16:24.920638Z"
    }
   },
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import torch.nn as nn\n",
    "import h5py\n",
    "import os\n",
    "from torch.utils.data import Dataset\n",
    "import torch.nn.functional as F"
   ],
   "outputs": [],
   "execution_count": 1
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "模型单元",
   "id": "936ebafb546c6818"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "class MultimodalDataset(Dataset):\n",
    "    def __init__(self, h5_path: str):\n",
    "        self.h5 = h5py.File(h5_path, 'r')\n",
    "        self.keys = list(self.h5.keys())\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.keys)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        group = self.h5[self.keys[idx]]\n",
    "        return {\n",
    "            \"images\": (\n",
    "                torch.from_numpy(group[\"image1\"][:]),\n",
    "                torch.from_numpy(group[\"image2\"][:])\n",
    "            ),\n",
    "            \"text\": torch.from_numpy(group[\"text\"][:]),\n",
    "            \"meta\": {\"folder\": group.attrs[\"folder\"]}\n",
    "        }"
   ],
   "id": "39078f1b26521dc4"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-25T02:16:44.409565Z",
     "start_time": "2025-03-25T02:16:44.400915Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class MLA(nn.Module):\n",
    "    def __init__(self, img_channels, text_dim=768, num_heads=8, dropout=0.1):\n",
    "        super().__init__()\n",
    "        # 确保维度对齐\n",
    "        assert text_dim % num_heads == 0, \"text_dim must be divisible by num_heads\"\n",
    "        \n",
    "        # 图像通道投影层（将图像维度与文本对齐）\n",
    "        self.img_proj = nn.Conv2d(img_channels, text_dim, kernel_size=1)\n",
    "        \n",
    "        # 多头注意力层\n",
    "        self.multihead_attn = nn.MultiheadAttention(\n",
    "            embed_dim=text_dim,\n",
    "            num_heads=num_heads,\n",
    "            dropout=dropout,\n",
    "            batch_first=True\n",
    "        )\n",
    "        \n",
    "        # 输出投影层（恢复图像原始通道）\n",
    "        self.output_proj = nn.Conv2d(text_dim, img_channels, kernel_size=1)\n",
    "        \n",
    "        # 层归一化\n",
    "        self.norm = nn.LayerNorm(text_dim)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, img, text):\n",
    "        \n",
    "        batch, channels, h, w = img.shape\n",
    "\n",
    "        img_projected = self.img_proj(img)  # (b, text_dim, h, w)\n",
    "\n",
    "        img_sequence = img_projected.flatten(2).permute(0, 2, 1)  # (b, h*w, text_dim)\n",
    "        \n",
    "        attn_output, _ = self.multihead_attn(\n",
    "            query=img_sequence,\n",
    "            key=text,\n",
    "            value=text\n",
    "        )\n",
    "\n",
    "        attn_output = self.norm(img_sequence + self.dropout(attn_output))\n",
    "\n",
    "        output_sequence = attn_output.permute(0, 2, 1)  # (b, text_dim, h*w)\n",
    "        output_projected = output_sequence.view(batch, 768, h, w)\n",
    "\n",
    "        output_img = self.output_proj(output_projected)  # b, c, h, w\n",
    "        \n",
    "        return output_img"
   ],
   "id": "51b60432f020d6af",
   "outputs": [],
   "execution_count": 2
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-25T02:16:44.988495Z",
     "start_time": "2025-03-25T02:16:44.976505Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class CNN_block_D(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels):\n",
    "        super(CNN_block_D, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)\n",
    "        self.batch = nn.BatchNorm2d(out_channels)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.con2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)\n",
    "        self.attn = MLA(img_channels=out_channels, num_heads=4)\n",
    "        self.pool = nn.MaxPool2d(2, 2)\n",
    "    \n",
    "    def forward(self, x, text):\n",
    "        x = self.conv1(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.con2(x)\n",
    "        x = self.batch(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.attn(x, text)\n",
    "        x = self.batch(x)\n",
    "        x = self.pool(x)\n",
    "        return x"
   ],
   "id": "80804df5b56102c3",
   "outputs": [],
   "execution_count": 3
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-25T02:16:45.518830Z",
     "start_time": "2025-03-25T02:16:45.506842Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class CNN_block_U(nn.Module):\n",
    "    def __init__(self, In_channels, Out_channels):\n",
    "        super().__init__()\n",
    "        # 通道数减半并尺寸倍增\n",
    "        self.conv_trans = nn.ConvTranspose2d(in_channels=In_channels,out_channels=Out_channels,kernel_size=3,stride=2,padding=1,output_padding=1)\n",
    "        self.bn = nn.BatchNorm2d(Out_channels)\n",
    "        self.relu = nn.SiLU()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv_trans(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.bn(x)\n",
    "        return x"
   ],
   "id": "4767627695175fcc",
   "outputs": [],
   "execution_count": 4
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-25T02:16:45.984505Z",
     "start_time": "2025-03-25T02:16:45.977324Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class CNN_block_33(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels):\n",
    "        super(CNN_block_33, self).__init__()\n",
    "        self.conv = nn.Sequential(\n",
    "            nn.Conv2d(in_channels, out_channels, kernel_size=33, padding=16),\n",
    "            nn.BatchNorm2d(out_channels)\n",
    "        )\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.conv(x)"
   ],
   "id": "af5026519f5fe696",
   "outputs": [],
   "execution_count": 5
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-25T02:16:46.527069Z",
     "start_time": "2025-03-25T02:16:46.519072Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class CNN_block_3(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels):\n",
    "        super(CNN_block_3, self).__init__()\n",
    "        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)\n",
    "        self.relu = nn.LeakyReLU()\n",
    "        self.norm = nn.BatchNorm2d(out_channels)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        return self.norm(self.relu(self.conv(x)))"
   ],
   "id": "1d18d4c3cd5ea3f6",
   "outputs": [],
   "execution_count": 6
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-25T02:16:47.005720Z",
     "start_time": "2025-03-25T02:16:46.993720Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class CNN_block_1(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels):\n",
    "        super(CNN_block_1, self).__init__()\n",
    "        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)\n",
    "        self.relu = nn.LeakyReLU()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.relu(self.conv(x))"
   ],
   "id": "5194dff42e29ea54",
   "outputs": [],
   "execution_count": 7
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-25T02:16:47.478817Z",
     "start_time": "2025-03-25T02:16:47.466808Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class atten_Gate(nn.Module):\n",
    "    def __init__(self,F_int):\n",
    "        super().__init__()\n",
    "        self.W_g = nn.Conv2d(F_int, F_int, 1)\n",
    "        self.W_x = nn.Conv2d(F_int, F_int, 1)\n",
    "        self.psi = nn.Conv2d(F_int, 1, 1)    # 注意力系数生成\n",
    "        self.relu = nn.LeakyReLU()\n",
    "        self.sigmoid = nn.Sigmoid()\n",
    "        self.norm = nn.BatchNorm2d(F_int)\n",
    "\n",
    "    def forward(self, g, x):\n",
    "        #print(g.shape)\n",
    "        g_trans = self.W_g(g)\n",
    "        x_trans = self.W_x(x)\n",
    "        \n",
    "        fused = self.relu(g_trans + x_trans)\n",
    "        alpha = self.sigmoid(self.psi(fused))\n",
    "\n",
    "        return self.norm(x * alpha)"
   ],
   "id": "890b230932b10589",
   "outputs": [],
   "execution_count": 8
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-25T02:16:49.148235Z",
     "start_time": "2025-03-25T02:16:47.980826Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class EdgeEnhancer(nn.Module):\n",
    "    def __init__(self, alpha=0.7, Device='cuda'):\n",
    "        super().__init__()\n",
    "        self.alpha = alpha  # 边缘增强强度系数\n",
    "        self.sobel_kernel = self._init_sobel_kernel().to(Device)\n",
    "\n",
    "    def _init_sobel_kernel(self):\n",
    "        # Sobel算子卷积核 [3,8]\n",
    "        kernel_x = torch.tensor([[1, 0, -1], \n",
    "                                [2, 0, -2],\n",
    "                                [1, 0, -1]], dtype=torch.float32).view(1,1,3,3)\n",
    "        \n",
    "        kernel_y = torch.tensor([[1, 2, 1],\n",
    "                                [0, 0, 0],\n",
    "                                [-1, -2, -1]], dtype=torch.float32).view(1,1,3,3)\n",
    "        return torch.cat([kernel_x, kernel_y], dim=0)  # shape(2,1,3,3)\n",
    "\n",
    "    def forward(self, x):\n",
    "        edges = F.conv2d(x, self.sobel_kernel, padding=1, groups=1)  # (n,2,512,512)\n",
    "        edges = torch.sqrt(edges[:,0]**2 + edges[:,1]**2)  \n",
    "        edges = edges.unsqueeze(1)\n",
    "        \n",
    "\n",
    "        enhanced = (1 - self.alpha) * x + self.alpha * edges\n",
    "        \n",
    "        return torch.clamp(enhanced, 0.0, 1.0)\n",
    "\n",
    "Edge = EdgeEnhancer()"
   ],
   "id": "8d2a76baf392c3d8",
   "outputs": [],
   "execution_count": 9
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-25T02:16:49.186356Z",
     "start_time": "2025-03-25T02:16:49.162362Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class UNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(UNet, self).__init__()\n",
    "        self.convD0 = CNN_block_33(1,16)\n",
    "        self.convF = CNN_block_3(16,64)\n",
    "        self.convD1 = CNN_block_D(64,128)\n",
    "        self.convD2 = CNN_block_D(128,256)\n",
    "        self.convD3 = CNN_block_D(256,512)\n",
    "        self.convD4 = CNN_block_D(512,1024)\n",
    "        \n",
    "        self.convU1 = CNN_block_U(1024,512)\n",
    "        self.convU2 = CNN_block_U(512,256)\n",
    "        self.convU3 = CNN_block_U(256,128)\n",
    "        self.convU4 = CNN_block_U(128,64)\n",
    "        \n",
    "        self.mixU = MLA(img_channels = 1024, num_heads = 8)\n",
    "        \n",
    "        self.jump1 = atten_Gate(64)\n",
    "        self.jump2 = atten_Gate(128)\n",
    "        self.jump3 = atten_Gate(256)\n",
    "        self.jump4 = atten_Gate(512)\n",
    "        \n",
    "        self.dec1 = CNN_block_3(64,16)\n",
    "        self.dec2 = CNN_block_3(16,4)\n",
    "        self.dec3 = CNN_block_1(4,1)\n",
    "        \n",
    "    def forward(self, img, text_emb):\n",
    "        x = self.convD0(img)\n",
    "        x = self.convF(x)   #c = 64\n",
    "        \n",
    "        mem1 = x\n",
    "        x = self.convD1(x,text_emb)\n",
    "        mem2 = x\n",
    "        x = self.convD2(x,text_emb)\n",
    "        mem3 = x\n",
    "        x = self.convD3(x,text_emb)\n",
    "        mem4 = x\n",
    "        x = self.convD4(x,text_emb)\n",
    "        x = self.mixU(x, text_emb)\n",
    "        \n",
    "        x = self.convU1(x)\n",
    "        x = self.jump4(mem4, x)\n",
    "        x = self.convU2(x)\n",
    "        x = self.jump3(mem3, x)\n",
    "        x = self.convU3(x)\n",
    "        x = self.jump2(mem2, x)\n",
    "        x = self.convU4(x)\n",
    "        x = self.jump1(mem1, x)\n",
    "        \n",
    "        x = self.dec1(x)\n",
    "        x = self.dec2(x)\n",
    "        return self.dec3(x)"
   ],
   "id": "6b9b3c8739968410",
   "outputs": [],
   "execution_count": 10
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-25T02:16:51.213266Z",
     "start_time": "2025-03-25T02:16:51.203119Z"
    }
   },
   "cell_type": "code",
   "source": "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')",
   "id": "ba077aae2386db83",
   "outputs": [],
   "execution_count": 11
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-25T02:17:19.895215Z",
     "start_time": "2025-03-25T02:17:14.431444Z"
    }
   },
   "cell_type": "code",
   "source": [
    "model = torch.load('D:/下载/model.pth')\n",
    "model.eval()"
   ],
   "id": "a78c55d096f0fab",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "UNet(\n",
       "  (convD0): CNN_block_33(\n",
       "    (conv): Sequential(\n",
       "      (0): Conv2d(1, 16, kernel_size=(33, 33), stride=(1, 1), padding=(16, 16))\n",
       "      (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (convF): CNN_block_3(\n",
       "    (conv): Conv2d(16, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (relu): ReLU()\n",
       "    (norm): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       "  (convD1): CNN_block_D(\n",
       "    (conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (batch): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (relu): ReLU()\n",
       "    (con2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (attn): MLA(\n",
       "      (img_proj): Conv2d(128, 768, kernel_size=(1, 1), stride=(1, 1))\n",
       "      (multihead_attn): MultiheadAttention(\n",
       "        (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)\n",
       "      )\n",
       "      (output_proj): Conv2d(768, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "      (norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (dropout): Dropout(p=0.1, inplace=False)\n",
       "    )\n",
       "    (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "  )\n",
       "  (convD2): CNN_block_D(\n",
       "    (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (batch): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (relu): ReLU()\n",
       "    (con2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (attn): MLA(\n",
       "      (img_proj): Conv2d(256, 768, kernel_size=(1, 1), stride=(1, 1))\n",
       "      (multihead_attn): MultiheadAttention(\n",
       "        (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)\n",
       "      )\n",
       "      (output_proj): Conv2d(768, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "      (norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (dropout): Dropout(p=0.1, inplace=False)\n",
       "    )\n",
       "    (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "  )\n",
       "  (convD3): CNN_block_D(\n",
       "    (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (batch): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (relu): ReLU()\n",
       "    (con2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (attn): MLA(\n",
       "      (img_proj): Conv2d(512, 768, kernel_size=(1, 1), stride=(1, 1))\n",
       "      (multihead_attn): MultiheadAttention(\n",
       "        (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)\n",
       "      )\n",
       "      (output_proj): Conv2d(768, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "      (norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (dropout): Dropout(p=0.1, inplace=False)\n",
       "    )\n",
       "    (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "  )\n",
       "  (convD4): CNN_block_D(\n",
       "    (conv1): Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (batch): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (relu): ReLU()\n",
       "    (con2): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (attn): MLA(\n",
       "      (img_proj): Conv2d(1024, 768, kernel_size=(1, 1), stride=(1, 1))\n",
       "      (multihead_attn): MultiheadAttention(\n",
       "        (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)\n",
       "      )\n",
       "      (output_proj): Conv2d(768, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "      (norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (dropout): Dropout(p=0.1, inplace=False)\n",
       "    )\n",
       "    (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "  )\n",
       "  (convU1): CNN_block_U(\n",
       "    (conv_trans): ConvTranspose2d(1024, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))\n",
       "    (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (relu): SiLU()\n",
       "  )\n",
       "  (convU2): CNN_block_U(\n",
       "    (conv_trans): ConvTranspose2d(512, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))\n",
       "    (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (relu): SiLU()\n",
       "  )\n",
       "  (convU3): CNN_block_U(\n",
       "    (conv_trans): ConvTranspose2d(256, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))\n",
       "    (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (relu): SiLU()\n",
       "  )\n",
       "  (convU4): CNN_block_U(\n",
       "    (conv_trans): ConvTranspose2d(128, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))\n",
       "    (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (relu): SiLU()\n",
       "  )\n",
       "  (mixU): MLA(\n",
       "    (img_proj): Conv2d(1024, 768, kernel_size=(1, 1), stride=(1, 1))\n",
       "    (multihead_attn): MultiheadAttention(\n",
       "      (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)\n",
       "    )\n",
       "    (output_proj): Conv2d(768, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "    (norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "    (dropout): Dropout(p=0.1, inplace=False)\n",
       "  )\n",
       "  (jump1): atten_Gate(\n",
       "    (W_g): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "    (W_x): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "    (psi): Conv2d(64, 1, kernel_size=(1, 1), stride=(1, 1))\n",
       "    (relu): ReLU(inplace=True)\n",
       "    (sigmoid): Sigmoid()\n",
       "    (norm): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       "  (jump2): atten_Gate(\n",
       "    (W_g): Conv2d(128, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "    (W_x): Conv2d(128, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "    (psi): Conv2d(128, 1, kernel_size=(1, 1), stride=(1, 1))\n",
       "    (relu): ReLU(inplace=True)\n",
       "    (sigmoid): Sigmoid()\n",
       "    (norm): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       "  (jump3): atten_Gate(\n",
       "    (W_g): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "    (W_x): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "    (psi): Conv2d(256, 1, kernel_size=(1, 1), stride=(1, 1))\n",
       "    (relu): ReLU(inplace=True)\n",
       "    (sigmoid): Sigmoid()\n",
       "    (norm): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       "  (jump4): atten_Gate(\n",
       "    (W_g): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "    (W_x): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "    (psi): Conv2d(512, 1, kernel_size=(1, 1), stride=(1, 1))\n",
       "    (relu): ReLU(inplace=True)\n",
       "    (sigmoid): Sigmoid()\n",
       "    (norm): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       "  (dec1): CNN_block_3(\n",
       "    (conv): Conv2d(64, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (relu): ReLU()\n",
       "    (norm): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       "  (dec2): CNN_block_3(\n",
       "    (conv): Conv2d(16, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (relu): ReLU()\n",
       "    (norm): BatchNorm2d(4, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       "  (dec3): CNN_block_1(\n",
       "    (conv): Conv2d(4, 1, kernel_size=(1, 1), stride=(1, 1))\n",
       "    (relu): ReLU()\n",
       "    (norm): BatchNorm2d(1, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 13
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-25T02:17:41.698556Z",
     "start_time": "2025-03-25T02:17:40.485873Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def show_t(tensor):\n",
    "    tensor = tensor\n",
    "    # 转换为numpy数组\n",
    "    np_array = tensor.detach().cpu().numpy()\n",
    "    \n",
    "    # 显示图像\n",
    "    plt.imshow(np_array, cmap='gray', vmin=0, vmax=1)\n",
    "    plt.axis('off')\n",
    "    plt.show()"
   ],
   "id": "503898f33355d0f8",
   "outputs": [],
   "execution_count": 17
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-25T02:17:42.364276Z",
     "start_time": "2025-03-25T02:17:42.334243Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def eval():\n",
    "    directory = 'C:/Users/Lenovo/Desktop/DC/ds'\n",
    "    norm = nn.BatchNorm2d(1)\n",
    "    \n",
    "    for filename in os.listdir(directory):\n",
    "        file_path = os.path.join(directory, filename)\n",
    "        #print(file_path)\n",
    "        dataset = MultimodalDataset(file_path)\n",
    "        for i in range(5):\n",
    "            dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True)\n",
    "            for item in dataloader:\n",
    "                Input = norm(item['images'][0].to(device))\n",
    "                tar = norm(item['images'][1].to(device))\n",
    "                text = norm(item['text'].to(device))\n",
    "                pre = model.forward(Edge(Input),text)\n",
    "                show_t(Edge(Input.unsqueeze(0)).squeeze(0).squeeze(0))\n",
    "                show_t(pre.squeeze(0).squeeze(0))\n",
    "                break\n",
    "            break\n",
    "        break\n",
    "                "
   ],
   "id": "a2f8be470d591a45",
   "outputs": [],
   "execution_count": 18
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "eval()",
   "id": "65501dead1738672"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
